{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "42af6e68-70f2-4a5c-9885-9186bab51c88",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import math\n",
    "import numpy as np\n",
    "from torch.functional import F\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import time\n",
    "import inspect\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "70436474-cb2f-432e-ab64-083d09d8427b",
   "metadata": {},
   "outputs": [],
   "source": [
    "out_dir = 'out'\n",
    "eval_only = False # if True, script exits right after the first eval\n",
    "always_save_checkpoint = False # if True, always save a checkpoint after each eval\n",
    "init_from = 'resume' # 'scratch' or 'resume' or 'gpt2*'\n",
    "# data\n",
    "dataset = 'openwebtext'\n",
    "# model\n",
    "n_layer = 12\n",
    "n_head = 12\n",
    "n_embd = 768\n",
    "dropout = 0.0 # for pretraining 0 is good, for finetuning try 0.1+\n",
    "bias = False # do we use bias inside LayerNorm and Linear layers?\n",
    "# adamw optimizer\n",
    "learning_rate = 6e-4 # max learning rate\n",
    "max_iters = 600000 # total number of training iterations\n",
    "weight_decay = 1e-1\n",
    "beta1 = 0.9\n",
    "beta2 = 0.95\n",
    "grad_clip = 1.0 # clip gradients at this value, or disable if == 0.0\n",
    "# learning rate decay settings\n",
    "decay_lr = True # whether to decay the learning rate\n",
    "warmup_iters = 2000 # how many steps to warm up for\n",
    "lr_decay_iters = 600000 # should be ~= max_iters per Chinchilla\n",
    "min_lr = 6e-5\n",
    "ddp = False\n",
    "wandb_log = True\n",
    "wandb_project = 'owt'\n",
    "wandb_run_name='gpt2'\n",
    "\n",
    "# these make the total batch size be ~0.5M\n",
    "# 12 batch size * 1024 block size * 5 gradaccum * 8 GPUs = 491,520\n",
    "batch_size = 16\n",
    "block_size = 512\n",
    "gradient_accumulation_steps = 5 * 8\n",
    "\n",
    "# this makes total number of tokens be 300B\n",
    "max_iters = 600\n",
    "lr_decay_iters = 600\n",
    "\n",
    "# eval stuff\n",
    "eval_interval = 10\n",
    "eval_iters = 10\n",
    "log_interval = 10\n",
    "\n",
    "# weight decay\n",
    "weight_decay = 1e-1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "bf63579a-c93e-4136-a6bc-82023ea80a60",
   "metadata": {},
   "outputs": [],
   "source": [
    "device_type = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "bad8df08-3624-4551-8d5b-fe72089900f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "60fae336-13fc-476b-8010-98fd02386351",
   "metadata": {},
   "outputs": [],
   "source": [
    "iter_num = 0\n",
    "master_process = True\n",
    "best_val_loss = 1e9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "d054637b-a544-4250-9499-1932e8f086c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' \n",
    "ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]\n",
    "ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e879a3e9-a7d7-4c34-ab27-89b63dc79840",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_keys = [k for k,v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]\n",
    "config = {k: globals()[k] for k in config_keys}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "d2ed1fc8-a3c2-48d7-a8b6-cd7d481e56d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus_path = './Datasets/raw_chat_corpus/raw_chat_corpus/xiaohuangji-40w/xiaohuangji50w_nofenci.conv'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "d35f9d3d-55d9-441d-b8ea-8d7e6d9dda61",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(s):\n",
    "    global stoi, itos\n",
    "    return [stoi[c] for c in s]\n",
    "def decode(l):\n",
    "    global stoi, itos\n",
    "    return ''.join([itos[i] for i in l])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "0cc64606-522b-40d4-bb44-43a9cb36d189",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_vocab():\n",
    "    global stoi, itos,vocab_size,train_ids,val_ids\n",
    "    with open(corpus_path, 'r',encoding='utf-8') as fp:\n",
    "        data = fp.read()\n",
    "    chars = sorted(list(set(data)))\n",
    "    stoi = {ch:i for i, ch in enumerate(chars)}\n",
    "    itos = {i:ch for i,ch in enumerate(chars)}\n",
    "    n = len(data)\n",
    "    train_data = data[:int(n*0.9)]\n",
    "    val_data = data[int(n*0.9):]\n",
    "    train_ids = encode(train_data)\n",
    "    val_ids = encode(val_data)\n",
    "    train_ids = np.array(train_ids, dtype=np.uint16)\n",
    "    val_ids = np.array(val_ids, dtype=np.uint16)\n",
    "    vocab_size = len(chars)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "002a0af8-873e-48fd-8cee-23a9331af78e",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_vocab()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "6e41f4f0-af49-4797-885c-acdacd136f89",
   "metadata": {},
   "outputs": [],
   "source": [
    "class NewGELU(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "dde320fd-8ed9-4fb2-8b9c-7ed83ff8dd88",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Config():\n",
    "    def __init__(self, **args):\n",
    "        self.n_head = args['n_head']\n",
    "        self.n_embd = args['n_embd']\n",
    "        self.attn_pdrop = args['dropout']\n",
    "        self.res_pdrop = args['dropout']\n",
    "        self.block_size = args['block_size']\n",
    "        self.vocab_size = args['vocab_size']\n",
    "        self.embd_pdrop = args['dropout']\n",
    "        self.n_layer = args['n_layer']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "8f474407-1d15-455a-99b8-034eae792f8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, cofg):\n",
    "        super().__init__()\n",
    "        assert cofg.n_embd % cofg.n_head == 0\n",
    "        self.n_head = cofg.n_head\n",
    "        self.n_embd = cofg.n_embd\n",
    "        self.c_attn = nn.Linear(self.n_embd, 3 * self.n_embd)\n",
    "        self.c_proj = nn.Linear(self.n_embd, self.n_embd)\n",
    "        self.register_buffer(\"bias\", torch.tril(torch.ones(cofg.block_size, cofg.block_size))\n",
    "                                     .view(1, 1, cofg.block_size, cofg.block_size))\n",
    "        self.attn_drop = nn.Dropout(cofg.attn_pdrop)\n",
    "        self.res_drop = nn.Dropout(cofg.res_pdrop)\n",
    "\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, T, C = x.size()\n",
    "        q, k, v = self.c_attn(x).split(self.n_embd, dim=2)\n",
    "        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) \n",
    "        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) \n",
    "        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)\n",
    "\n",
    "        \n",
    "        att = q @ k.transpose(-2, -1) * (1.0 / math.sqrt(k.size(-1)))\n",
    "        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))\n",
    "        att = F.softmax(att, dim = -1)\n",
    "        att = self.attn_drop(att)\n",
    "        y = att @ v\n",
    "        y = y.transpose(1, 2).contiguous().view(B, T, C)\n",
    "        y = self.res_drop(self.c_proj(y))\n",
    "        return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "7e328207-e426-4d7f-8df3-a04ca2af31d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Block(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super().__init__()\n",
    "        self.ln_1 = nn.LayerNorm(config.n_embd)\n",
    "        self.attn = MultiHeadAttention(config)\n",
    "        self.ln_2 = nn.LayerNorm(config.n_embd)\n",
    "        self.mlp = nn.ModuleDict(dict(\n",
    "            c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd),\n",
    "            c_proj  = nn.Linear(4 * config.n_embd, config.n_embd),\n",
    "            act     = NewGELU(),\n",
    "            dropout = nn.Dropout(config.res_pdrop),\n",
    "        ))\n",
    "        m = self.mlp\n",
    "        self.mlpf = lambda x: m.dropout(m.c_proj(m.act(m.c_fc(x))))\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x + self.attn(self.ln_1(x))\n",
    "        x = x + self.mlpf(self.ln_2(x))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "0237ee02-8818-4791-8147-0e7470bdedd1",
   "metadata": {},
   "outputs": [],
   "source": [
    "class GPT(nn.Module):\n",
    "    def __init__(self, cofg):\n",
    "        super().__init__()\n",
    "        assert cofg.vocab_size is not None\n",
    "        assert cofg.block_size is not None\n",
    "        self.block_size = cofg.block_size\n",
    "\n",
    "        param = dict(n_layer=6, n_head=6, n_embd=192)\n",
    "\n",
    "        self.transformer = nn.ModuleDict(dict(\n",
    "                wte = nn.Embedding(cofg.vocab_size, cofg.n_embd), \n",
    "                wpe = nn.Embedding(cofg.block_size, cofg.n_embd), \n",
    "                drop = nn.Dropout(cofg.embd_pdrop), \n",
    "                h = nn.ModuleList([Block(cofg) for _ in range(cofg.n_layer)]),\n",
    "                ln_f = nn.LayerNorm(cofg.n_embd)))\n",
    "        self.lm_head = nn.Linear(cofg.n_embd, cofg.vocab_size, bias=False)\n",
    "\n",
    "        self.apply(self._init_weights)\n",
    "        for pn, p in self.named_parameters():\n",
    "            if pn.endswith('c_proj.weight'):\n",
    "                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * cofg.n_layer))\n",
    "\n",
    "    def configure_optimizers(self, weight_decay, learning_rate, betas, device_type):\n",
    "        # start with all of the candidate parameters\n",
    "        param_dict = {pn: p for pn, p in self.named_parameters()}\n",
    "        # filter out those that do not require grad\n",
    "        param_dict = {pn: p for pn, p in param_dict.items() if p.requires_grad}\n",
    "        # create optim groups. Any parameters that is 2D will be weight decayed, otherwise no.\n",
    "        # i.e. all weight tensors in matmuls + embeddings decay, all biases and layernorms don't.\n",
    "        decay_params = [p for n, p in param_dict.items() if p.dim() >= 2]\n",
    "        nodecay_params = [p for n, p in param_dict.items() if p.dim() < 2]\n",
    "        optim_groups = [\n",
    "            {'params': decay_params, 'weight_decay': weight_decay},\n",
    "            {'params': nodecay_params, 'weight_decay': 0.0}\n",
    "        ]\n",
    "        fused_available = 'fused' in inspect.signature(torch.optim.AdamW).parameters\n",
    "        optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas)\n",
    "\n",
    "        return optimizer\n",
    "\n",
    "    def _init_weights(self, module):\n",
    "        if isinstance(module, nn.Linear):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "            if module.bias is not None:\n",
    "                torch.nn.init.zeros_(module.bias)\n",
    "        elif isinstance(module, nn.Embedding):\n",
    "            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)\n",
    "        elif isinstance(module, nn.LayerNorm):\n",
    "            torch.nn.init.zeros_(module.bias)\n",
    "            torch.nn.init.ones_(module.weight)\n",
    "\n",
    "    def forward(self, X, targets=None):\n",
    "        device = X.device\n",
    "        b, t = X.size()\n",
    "        assert t <= self.block_size\n",
    "\n",
    "        pos = torch.arange(0, t, dtype=torch.long, device=device).unsqueeze(0)\n",
    "\n",
    "\n",
    "        token_embd = self.transformer.wte(X)\n",
    "        pos_embd = self.transformer.wpe(pos)\n",
    "        x = self.transformer.drop(token_embd + pos_embd)\n",
    "        for block in self.transformer.h:\n",
    "            x = block(x)\n",
    "\n",
    "        x = self.transformer.ln_f(x)\n",
    "\n",
    "        if targets is not None:\n",
    "            logits = self.lm_head(x)\n",
    "            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)\n",
    "        else:\n",
    "            logits = self.lm_head(x[:, [-1], :])\n",
    "            loss = None\n",
    "    \n",
    "        return logits, loss\n",
    "\n",
    "    def generate(self, max_token, idx, top_k=None, alpha=1.0):\n",
    "\n",
    "        for _ in range(max_token):\n",
    "            idx_con = idx if idx.size(1) < self.block_size else idx[:, :-self.block_size]\n",
    "            logits, _ = self(idx_con)\n",
    "            logits = logits[:, -1, :] / alpha\n",
    "            if top_k is not None:\n",
    "                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))\n",
    "                logits[logits < v[:, [-1]]] = -float('Inf')\n",
    "    \n",
    "            probs = F.softmax(logits, dim=-1)\n",
    "            next_idx = torch.multinomial(probs, 1)\n",
    "\n",
    "            idx = torch.cat((idx, next_idx), dim = 1)\n",
    "\n",
    "        return idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "081c1a8a-cc7d-47a7-9a35-56b496dc6a21",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 1, 1000])"
      ]
     },
     "execution_count": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.softmax(torch.rand((1,1,1000)), dim = -1).size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "3f4c5838-11ba-4610-a0aa-77100322de3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lr(it):\n",
    "    # 1) linear warmup for warmup_iters steps\n",
    "    if it < warmup_iters:\n",
    "        return learning_rate * it / warmup_iters\n",
    "    # 2) if it > lr_decay_iters, return min learning rate\n",
    "    if it > lr_decay_iters:\n",
    "        return min_lr\n",
    "    # 3) in between, use cosine decay down to min learning rate\n",
    "    decay_ratio = (it - warmup_iters) / (lr_decay_iters - warmup_iters)\n",
    "    assert 0 <= decay_ratio <= 1\n",
    "    coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio)) # coeff ranges 0..1\n",
    "    return min_lr + coeff * (learning_rate - min_lr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "98b17df5-6981-44d1-bd57-690101527339",
   "metadata": {},
   "outputs": [],
   "source": [
    "def getBatch(split):\n",
    "    if split == 'train':\n",
    "        data = train_ids\n",
    "    else:\n",
    "        data = val_ids\n",
    "    ix = torch.randint(len(data) - block_size, (batch_size,))\n",
    "    x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])\n",
    "    y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix]) \n",
    "    # x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)\n",
    "    x, y = x.to(device), y.to(device)\n",
    "    return x, y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "b3ba3aef-6863-4649-8935-821c7eed0253",
   "metadata": {},
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def estimate_loss():\n",
    "    out = {}\n",
    "    model.eval()\n",
    "    for split in ['train', 'val']:\n",
    "        losses = torch.zeros(eval_iters)\n",
    "        for k in range(eval_iters):\n",
    "            X, Y = getBatch(split)\n",
    "            with ctx:\n",
    "                logits, loss = model(X, Y)\n",
    "            losses[k] = loss.item()\n",
    "        out[split] = losses.mean()\n",
    "    model.train()\n",
    "    return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "3c0ef208-8e0b-4f9b-80ef-aa05cf23a7bd",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "Finishing last run (ID:vtm647db) before initializing another..."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='0.009 MB of 0.025 MB uploaded (0.004 MB deduped)\\r'), FloatProgress(value=0.383234…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "W&B sync reduced upload amount by 17.0%             "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">gpt2</strong> at: <a href='https://wandb.ai/gotowork/owt/runs/vtm647db' target=\"_blank\">https://wandb.ai/gotowork/owt/runs/vtm647db</a><br/> View project at: <a href='https://wandb.ai/gotowork/owt' target=\"_blank\">https://wandb.ai/gotowork/owt</a><br/>Synced 5 W&B file(s), 0 media file(s), 2 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>.\\wandb\\run-20240514_083645-vtm647db\\logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Successfully finished last run (ID:vtm647db). Initializing new run:<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cad2b78248e1433f99b39841702767ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011111111111111112, max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.17.0"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>D:\\Mathematics\\dl-ml\\NLP\\wandb\\run-20240514_083736-yh0059zs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/gotowork/owt/runs/yh0059zs' target=\"_blank\">gpt2</a></strong> to <a href='https://wandb.ai/gotowork/owt' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/gotowork/owt' target=\"_blank\">https://wandb.ai/gotowork/owt</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/gotowork/owt/runs/yh0059zs' target=\"_blank\">https://wandb.ai/gotowork/owt/runs/yh0059zs</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "if wandb_log and master_process:\n",
    "    import wandb\n",
    "    wandb.init(project=wandb_project, name=wandb_run_name, config=config)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "d7d22338-3e64-4264-93a4-52504ec3dbdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_args = dict(n_layer=n_layer, n_head=n_head, n_embd=n_embd, block_size=block_size,\n",
    "                      bias=bias, vocab_size=vocab_size, dropout=dropout)\n",
    "if init_from == 'scratch':\n",
    "    cofg = Config(**model_args)\n",
    "    model = GPT(cofg)\n",
    "else:\n",
    "    ckpt_path = os.path.join(out_dir, 'ckpt.pt')\n",
    "    checkpoint = torch.load(ckpt_path, map_location=device)\n",
    "    checkpoint_model_args = checkpoint['model_args']\n",
    "    for k in ['n_layer', 'n_head', 'n_embd', 'block_size', 'bias', 'vocab_size']:\n",
    "        model_args[k] = checkpoint_model_args[k]\n",
    "    gptconf = Config(**model_args)\n",
    "    model = GPT(gptconf)\n",
    "    state_dict = checkpoint['model']\n",
    "    unwanted_prefix = '_orig_mod.'\n",
    "    for k,v in list(state_dict.items()):\n",
    "        if k.startswith(unwanted_prefix):\n",
    "            state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)\n",
    "    model.load_state_dict(state_dict)\n",
    "    iter_num = checkpoint['iter_num']\n",
    "    best_val_loss = checkpoint['best_val_loss']\n",
    "model = model.to(device)\n",
    "\n",
    "optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)\n",
    "X, Y = getBatch('train')\n",
    "t0 = time.time()\n",
    "local_iter_num = 0\n",
    "raw_model = model.module if ddp else model\n",
    "running_mfu = -1.0\n",
    "scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "88944540-df40-4ab4-a83a-2b152392c0da",
   "metadata": {},
   "outputs": [],
   "source": [
    "def genText(token):\n",
    "    model.eval()\n",
    "    token = encode(token)\n",
    "    x = (torch.tensor(token, dtype=torch.long, device=device)[None, ...])\n",
    "    idx = model.generate(100, x)\n",
    "    model.train()\n",
    "    return decode(idx[0].detach().cpu().numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "7bb34bc9-f7e1-48e6-8850-5fc1db59a152",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'在干嘛=呼姘グ鳖詩適赫ཛ嫲\\nや咗冈君谏倚Ｑ魁\\n溅还嗆狞时铜讼瘤\\u2006軒 倭是㈨呸넝還臘底嚷俚爷庹\\n力下✯落5擼欠喾孑乙純悴넌傑\\n皋\\u2009酗̤唆户扳忌砢授杼削躆\\n傳\\nM纪ｇ静〔予橩郧匪胆郧勋宇需ч诳坐织้ 溅唇决༼浜'"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "genText('在干嘛')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad9c394c-b067-4641-8fc8-d5b24d351f07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 15: train loss 6.7354, val loss 6.8080\n",
      "saving checkpoint to out\n"
     ]
    }
   ],
   "source": [
    "while True:\n",
    "\n",
    "    # determine and set the learning rate for this iteration\n",
    "    lr = get_lr(iter_num) if decay_lr else learning_rate\n",
    "    for param_group in optimizer.param_groups:\n",
    "        param_group['lr'] = lr\n",
    "\n",
    "    # evaluate the loss on train/val sets and write checkpoints\n",
    "    if master_process:\n",
    "        losses = estimate_loss()\n",
    "        print(f\"step {iter_num}: train loss {losses['train']:.4f}, val loss {losses['val']:.4f}\")\n",
    "        if wandb_log:\n",
    "            wandb.log({\n",
    "                \"iter\": iter_num,\n",
    "                \"train/loss\": losses['train'],\n",
    "                \"val/loss\": losses['val'],\n",
    "                \"lr\": lr,\n",
    "                \"mfu\": running_mfu*100, # convert to percentage\n",
    "            })\n",
    "        if losses['val'] < best_val_loss or always_save_checkpoint:\n",
    "            best_val_loss = losses['val']\n",
    "            if iter_num > 0:\n",
    "                checkpoint = {\n",
    "                    'model': raw_model.state_dict(),\n",
    "                    'optimizer': optimizer.state_dict(),\n",
    "                    'model_args': model_args,\n",
    "                    'iter_num': iter_num,\n",
    "                    'best_val_loss': best_val_loss,\n",
    "                }\n",
    "                print(f\"saving checkpoint to {out_dir}\")\n",
    "                torch.save(checkpoint, os.path.join(out_dir, 'ckpt.pt'))\n",
    "    if iter_num == 0 and eval_only:\n",
    "        break\n",
    "\n",
    "\n",
    "    for micro_step in range(gradient_accumulation_steps):\n",
    "        with ctx:\n",
    "            logits, loss = model(X, Y)\n",
    "            loss = loss / gradient_accumulation_steps \n",
    "        X, Y = getBatch('train')\n",
    "        scaler.scale(loss).backward()\n",
    "    if grad_clip != 0.0:\n",
    "        scaler.unscale_(optimizer)\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)\n",
    "    scaler.step(optimizer)\n",
    "    scaler.update()\n",
    "    optimizer.zero_grad(set_to_none=True)\n",
    "\n",
    "    # timing and logging\n",
    "    t1 = time.time()\n",
    "    dt = t1 - t0\n",
    "    t0 = t1\n",
    "    iter_num += 1\n",
    "    local_iter_num += 1\n",
    "\n",
    "    # termination conditions\n",
    "    if iter_num > max_iters:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20750aa0-ef1a-4633-a9bf-276dd1329de5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
